"""测试各种识别算法的统一接口"""

import numpy as np
import cv2
import tqdm
import argparse
from utils.ellipse import from_ellipse_args, ellipse_fit, ellipse_args

# 识别数据库
from database.plain.binad import XuBuilder, ChenBuilder, ForsythBuilder
from database.plain.triad import HanakBuilder, ParkBuilder, ChristianBuilder
from database.hash import XuHashBuilder, ChristianHashBuilder
from database.hierarchy.binad import ForsythHierarchyBuilder, XuHierarchyBuilder
from database.hierarchy.binad import ChenHierarchyBuilder
from database.hierarchy.triad import ChristianHierarchyBuilder, ParkHierarchyBuilder

# 识别方法
from match.pyramid.binad_fast import XuFastPyramid, ChenFastPyramid, ForsythFastPyramid
from match.pyramid.triad import ParkPyramid, ChristianPyramid
from match.pyramid.triad_vote import ParkPyramidVote, ChristianPyramidVote
from match.hash.binad_vote import XuHashVote
from match.hash.triad import ChristianHash
from match.hash.triad_vote import ChristianHashVote

method_list = {
    # "Xu": (XuBuilder, XuPyramid),  # deprecated
    # "Chen": (ChenBuilder, ChenPyramid), # deprecated
    # binad plain method
    "XuFast": (XuBuilder, XuFastPyramid),
    # "ChenFast": (ChenBuilder, ChenFastPyramid),
    # "ForsythFast": (ForsythBuilder, ForsythFastPyramid),
    # # triad plain method
    # # "ParkFast": (ParkBuilder, ParkPyramid),
    # # "Christian": (ChristianBuilder, ChristianPyramid),
    # # Hash method
    # # "XuHash": (XuHashBuilder, XuHashVote),
    # # "ChristianHash": (ChristianHashBuilder, ChristianHash),
    # # binad hierarchy method
    # "ChenHierarchy": (ChenHierarchyBuilder, ChenFastPyramid),
    # "ForsythHierarchy": (ForsythHierarchyBuilder, ForsythFastPyramid),
    # "XuHierarchy": (XuHierarchyBuilder, XuFastPyramid),
    # # triad hierarchy method
    # "ChristianHierarchy": (ChristianHierarchyBuilder, ChristianPyramid),
    # "ParkHierarchy": (ParkHierarchyBuilder, ParkPyramid),
    # # triad vote
    # "ChristianVote": (ChristianBuilder, ChristianPyramidVote),
    # "ParkVote": (ParkBuilder, ParkPyramidVote),
}


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("--num", type=str, default=30, help="The number of points")
    arg_parser.add_argument("-f", type=float, default=25)
    arg_parser.add_argument("-d", type=float, default=5.5)
    arg_parser.add_argument("--reso", type=int, default=1024)
    arg_parser.add_argument("--noise", type=float, default=0)
    arg_parser.add_argument("--drift", type=float, default=0)
    arg_parser.add_argument("--height", type=float, default=4)
    arg_parser.add_argument("--rotate", type=float, default=np.pi / 5)
    arg_parser.add_argument(
        "--exp", type=int, default=100, help="The number of experiments"
    )
    arg_parser.add_argument("--recall", type=float, default=90)
    arg_parser.add_argument("--outlier", type=float, default=0)
    arg_parser.add_argument(
        "--use_init", action="store_true", help="Use pnp pose as initial"
    )
    arg_parser.add_argument(
        "--reproduce", action="store_true", help="Use default random seed"
    )
    arg_parser.add_argument("--debug", action="store_true", help="Debug mode")
    arg_parser.add_argument("--catalog", type=str, default="craters.csv")
    arg_parser.add_argument(
        "--pose",
        type=str,
        default="/home/a804_cbf/Code/luner_crater/navigator/database/camera_transforms.csv",
    )
    arg_parser.add_argument("--device", type=str, default="cuda:0")
    return arg_parser.parse_args()


def debug_triad(module, i, j, k, Qs, C, Qs_prime_idx, idx):
    ix, invariants = module.invariant(i, j, k, Qs[i], Qs[j], Qs[k, None])
    _, real_invariants = module.invariant(
        i, j, k, C[Qs_prime_idx[i]], C[Qs_prime_idx[j]], C[Qs_prime_idx[k], None]
    )
    ind = np.argwhere(
        (idx[0] == Qs_prime_idx[ix[0]])
        & (idx[1] == Qs_prime_idx[ix[1]])
        & (idx[2] == Qs_prime_idx[ix[2]])
    ).squeeze()
    identified_idx = idx[:, ind]
    real_idx = Qs_prime_idx[ix].squeeze()
    return invariants, real_invariants, identified_idx, real_idx


def debug_binad(module, i, j, Qs, C, Qs_prime_idx, idx):
    invariants = module.invariant(Qs[i, None], Qs[j, None])
    real_invariants = module.invariant(
        C[Qs_prime_idx[i], None], C[Qs_prime_idx[j], None]
    )
    ind = np.argwhere(
        (idx[0] == Qs_prime_idx[i]) & (idx[1] == Qs_prime_idx[j])
    ).squeeze()
    identified_idx = idx[:, ind]
    # identified_idx = identifier.idx[:,real_invariants]
    real_idx = Qs_prime_idx[[i, j]].squeeze()
    return invariants, real_invariants, identified_idx, real_idx


if __name__ == "__main__":
    args = arg_parse()
    if args.reproduce:
        np.random.seed(114514)
    for k, (BuilderClass, IdentifierClass) in method_list.items():
        name = BuilderClass.__name__
        builder = BuilderClass(
            args.catalog,
            max_field=1000,
            cache_path=f"./__pycache__/{name}.npz",
        )
        catalog_dict = builder()
        catalog = catalog_dict["catalog"]
        descriptor = catalog_dict["descriptor"]
        temp = catalog_dict["temp"]
        idx = catalog_dict["idx"]
        constants = catalog_dict["constants"]
        # 初始化识别器
        identifier = IdentifierClass(f"./__pycache__/{name}.npz", device=args.device)
        method_list[k] = identifier

    # 生成虚拟相机矩阵
    K = np.array(
        [
            [args.f / args.d * 1e3, 0, args.reso / 2],
            [0, args.f / args.d * 1e3, args.reso / 2],
            [0, 0, 1],
        ]
    )
    # 平面上任取三个点，构成生成子空间
    M = np.block([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
    # 此处不再是随意选取陨石坑圆，而是从库中读取
    elli_3d = []
    C = []
    cnt_3d = []
    # 生成椭圆的二维点
    for x, y, r in temp:
        points_2d = from_ellipse_args(
            x[None], y[None], r[None], r[None], np.zeros_like(x[None])
        )
        points_2d = np.pad(
            points_2d, ((0, 0), (0, 1), (0, 0)), "constant", constant_values=1
        )
        # 二维点转三维点，以模拟相机成像过程
        points_3d = np.einsum("ijk,jl->ilk", points_2d, M)
        C.append(np.array([[1, 0, -x], [0, 1, -y], [-x, -y, x * x + y * y - r * r]]))
        elli_3d.append(points_3d)
        cnt_3d.append((x, y, 0))
    C = np.array(C)
    cnt_3d = np.array(cnt_3d)
    elli_3d = np.concatenate(elli_3d, axis=0)
    # 初始化统计量
    output = {
        name: {"recall": [], "time": [], "failed": 0} for name in method_list.keys()
    }
    total = 0
    # with open(args.pose, "r") as f:
    #     f.readline()
    #     for line in tqdm.tqdm(f, desc="Testing Catalog"):
    gbar = tqdm.tqdm(total=args.exp, desc="Testing Catalog")
    while total < args.exp:
        # 生成像平面上的椭圆
        ## 相机随机位姿
        # quat = list(float(x) for x in line.strip().split(",")[4:])
        # R = Rotation.from_quat(quat).as_matrix()
        # Rvec = cv2.Rodrigues(R)[0]
        # position = list(float(x) for x in line.strip().split(",")[1:4])
        Rvec = np.random.randn(3, 1) * args.rotate
        R = cv2.Rodrigues(Rvec)[0]
        position = temp.mean(axis=0) + np.random.randn(3) * temp.std(axis=0)
        position[2] += args.height * 1e3
        T = -R @ position
        # P = K @ np.hstack((R, T[:, None]))
        # ## 生成像平面上的椭圆
        # H = P @ M.T
        # C_real = np.linalg.inv(H).T @ C @ np.linalg.inv(H)
        ## 在原始平面上采点，投影至像平面，再拟合为椭圆
        Qs = []
        Qs_args = []
        Us = []
        Qs_prime_idx = []
        obse_Qs_points = []
        points_2d = cv2.projectPoints(cnt_3d.T, Rvec, T, K.astype(np.float32), None)[
            0
        ].squeeze()
        ind = (
            (0 <= points_2d[:, 0])
            & (points_2d[:, 0] < args.reso)
            & (0 <= points_2d[:, 1])
            & (points_2d[:, 1] < args.reso)
        )
        for it, points_3d in enumerate(elli_3d):
            if not ind[it]:
                continue
            points_2d = cv2.projectPoints(
                points_3d[:3].T, Rvec, T, K.astype(np.float32), None
            )[0].squeeze()
            # 加入噪声和漂移
            points_2d += np.random.randn(*points_2d.shape) * args.noise + args.drift
            # 加入异常值点，即误检测的椭圆，同时删除一个正确的椭圆
            if np.random.rand() < args.outlier:
                r = np.random.uniform(10, 100, 2)
                theta = np.random.uniform(0, 2 * np.pi, 1)
                E = from_ellipse_args(
                    *np.block([np.random.uniform(0, 1024, 2), r, theta])
                )
                E, U = ellipse_fit(E[0], E[1])
            else:
                E, U = ellipse_fit(points_2d[:, 0][None], points_2d[:, 1][None])
                # C_real = cone_plane_cross(P.T @ E @ P, (0, 0, 1, 0))
            # 用当前椭圆的参数排除在视场以外的椭圆
            x, y, a, b, theta = ellipse_args(E)
            if x:
                if 0 <= x < args.reso and 0 <= y < args.reso:
                    Qs.append(E)
                    Us.append(U)
                    Qs_prime_idx.append(it)
                    obse_Qs_points.append(points_2d)
        # 考虑真实情况，同一幅图像中椭圆数量少于4个一定测不出位姿，大于20个不符合真实导航场景
        if len(Qs) < 4:
            continue
        # key = len(Qs)
        Qs = np.concatenate(Qs, axis=0)
        Us = np.concatenate(Us, axis=0)
        Qs_prime_idx = np.array(Qs_prime_idx)
        # # 测试图像平面上的椭圆不变量
        invariants, real_invariants, identified_idx, real_idx = debug_binad(
            builder, 0, 1, Qs, C, Qs_prime_idx, idx
        )
        # 对于Park方法的调试
        # invariants, real_invariants, identified_idx, real_idx = debug_triad(
        #     identifier, 0, 1, 2, Qs, C, Qs_prime_idx, idx
        # )
        # 再将全部陨石坑写入至文件，作为测试数据
        # with open(f"./output/30/arguments.csv", "a") as f:
        #     print(
        #         f"{args.f},{args.d},{args.reso}",
        #         ",".join(map(str, Rvec.squeeze())),
        #         ",".join(map(str, T.squeeze())),
        #         file=f,
        #         sep=",",
        #     )
        # np.savez(f"./output/30/{total}.npz", idx=Qs_prime_idx, points=obse_Qs_points)
        # 测试识别结果
        for method, identifier in method_list.items():
            try:
                t, result = identifier(
                    Qs,
                    Us,
                    factor=0.3,
                    K=K,
                    confidence=0.8,
                    th=3,
                    max_iter=10,
                    max_try=10,
                )
            except:
                output[method]["failed"] += 1
                break
            output[method]["time"].append(t)
            if result is None:
                output[method]["failed"] += 1
                continue
            # GT值 : Qs_prime_idx
            # 计算识别准确率
            output[method]["recall"].append(
                (result[1] == Qs_prime_idx[result[0]]).mean()
            )
        gbar.update(1)
        total += 1
    # 输出结果
    for method, identifier in method_list.items():
        with open(f"results/{method}.csv", "a") as w:
            w.write(
                f"{args.noise},{args.drift},{args.outlier},{args.height},{args.exp},{args.rotate:.4f},"
            )
            total_recall = np.mean(output[method]["recall"])
            total_time = np.mean(output[method]["time"])
            total_failed = output[method]["failed"]
            w.write(
                f"{1-total_failed/args.exp:.3f},{total_time:.3f},{total_recall:.3f}\n"
            )
